/* Copyright Thanh-Trung Pham, JAIST, 2013 */
package jp.ac.jaist.model.markov;

import org.apache.commons.collections.map.MultiKeyMap;
import org.apache.commons.math.linear.Array2DRowRealMatrix;
import org.apache.commons.math.linear.LUDecompositionImpl;
import org.apache.commons.math.linear.MatrixUtils;
import org.apache.commons.math.linear.RealMatrix;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

/**
 * @author Thanh-Trung Pham
 */
public class RMPIMarkovModel {
    private MultiKeyMap m_oTransientStates;
    private MultiKeyMap m_oAbsorbingStates;
    private double[][] m_oBMatrix;

    /**
     * The constructor.
     */
    public RMPIMarkovModel() {
        m_oTransientStates = new MultiKeyMap();
        m_oAbsorbingStates = new MultiKeyMap();
    }

    /**
     * Get a state by a text part and a set part.
     *
     * @param textPart    the text part.
     * @param setPart     the set part.
     * @param isAbsorbing true if the state is absorbing, otherwise, false.
     * @return the state.
     */
    public RMPIMarkovState getState(String textPart, Set<String> setPart, boolean isAbsorbing) {
        if (isAbsorbing) {
            return (RMPIMarkovState) m_oAbsorbingStates.get(textPart, setPart);
        } else {
            return (RMPIMarkovState) m_oTransientStates.get(textPart, setPart);
        }
    }

    /**
     * Add a new state.
     *
     * @param textPart    the text part.
     * @param setPart     the set part.
     * @param isAbsorbing true if the state is absorbing, otherwise, false.
     * @return the added state.
     */
    public RMPIMarkovState addState(String textPart, Set<String> setPart, boolean isAbsorbing) {
        RMPIMarkovState oState = new RMPIMarkovState(textPart, setPart, isAbsorbing);
        if (isAbsorbing) {
            oState.setID(m_oAbsorbingStates.size());
            m_oAbsorbingStates.put(textPart, setPart, oState);
        } else {
            oState.setID(m_oTransientStates.size());
            m_oTransientStates.put(textPart, setPart, oState);
        }
        return oState;
    }

    /**
     * Add a new transition.
     *
     * @param sourceTextPart        the source text part.
     * @param sourceSetPart         the source set part.
     * @param destinationTextPart   the destination text part.
     * @param destinationSetPart    the destination set part.
     * @param isAbsorbing           true if the state is absorbing, otherwise, false.
     * @param transitionProbability the transition probability.
     */
    public void addTransition(String sourceTextPart, Set<String> sourceSetPart, String destinationTextPart, Set<String> destinationSetPart, boolean isAbsorbing, double transitionProbability) {
        RMPIMarkovState oSourceState = getState(sourceTextPart, sourceSetPart, false);
        RMPIMarkovState oDestinationState = getState(destinationTextPart, destinationSetPart, isAbsorbing);

        addTransition(oSourceState, oDestinationState, transitionProbability);
    }

    /**
     * Add a new transition.
     *
     * @param sourceState           the source state.
     * @param destinationTextPart   the destination text part.
     * @param destinationSetPart    the destination set part.
     * @param isAbsorbing           true if the state is absorbing, otherwise, false.
     * @param transitionProbability the transition probability.
     */
    public void addTransition(RMPIMarkovState sourceState, String destinationTextPart, Set<String> destinationSetPart, boolean isAbsorbing, double transitionProbability) {
        RMPIMarkovState oDestinationState = getState(destinationTextPart, destinationSetPart, isAbsorbing);

        addTransition(sourceState, oDestinationState, transitionProbability);
    }

    /**
     * Add a new transition.
     *
     * @param sourceState           the source state.
     * @param destinationState      the destination state.
     * @param transitionProbability the transition probability.
     */
    public void addTransition(RMPIMarkovState sourceState, RMPIMarkovState destinationState, double transitionProbability) {
        RMPIMarkovTransition oTransition = new RMPIMarkovTransition(sourceState, destinationState, transitionProbability);

        sourceState.addOutgoingTransition(oTransition);
    }

    /**
     * Add a new transition.
     *
     * @param sourceTextPart        the source text part.
     * @param sourceSetPart         the source set part.
     * @param destinationState      the destination state.
     * @param transitionProbability the transition probability.
     */
    public void addTransition(String sourceTextPart, Set<String> sourceSetPart, RMPIMarkovState destinationState, double transitionProbability) {
        RMPIMarkovState oSourceState = getState(sourceTextPart, sourceSetPart, false);
        RMPIMarkovTransition oTransition = new RMPIMarkovTransition(oSourceState, destinationState, transitionProbability);

        oSourceState.addOutgoingTransition(oTransition);
    }

    /**
     * Create all-zero matrix.
     *
     * @param rowSize the row size.
     * @param colSize the column size.
     * @return the all-zero matrix.
     */
    private double[][] createAllZeroMatrix(int rowSize, int colSize) {
        double oMatrix[][] = new double[rowSize][colSize];
        Arrays.fill(oMatrix[0], 0.0);
        for (int i = 1; i < rowSize; i++) {
            System.arraycopy(oMatrix[0], 0, oMatrix[i], 0, colSize);
        }

        return oMatrix;
    }

    /**
     * Generate and then solve the Markov chain.
     */
    private void generateAndSolveMarkovChain() {
        if (m_oBMatrix == null) {
            int iRowSize = m_oTransientStates.size();

            double[][] oQMatrix = createAllZeroMatrix(iRowSize, iRowSize);
            double[][] oRMatrix = createAllZeroMatrix(iRowSize, m_oAbsorbingStates.size());

            for (Object oObj : m_oTransientStates.values()) {
                RMPIMarkovState oTransientState = (RMPIMarkovState) oObj;

                for (RMPIMarkovTransition oTransition : oTransientState.getOutgoingTransitions()) {
                    RMPIMarkovState oDesState = oTransition.getDestinationState();
                    if (oDesState.isAbsorbing()) {
                        oRMatrix[oTransition.getSourceState().getID()][oDesState.getID()] = oTransition.getTransitionProbability();
                    } else {
                        oQMatrix[oTransition.getSourceState().getID()][oDesState.getID()] = oTransition.getTransitionProbability();
                    }
                }
            }

            RealMatrix oRealQMatrix = new Array2DRowRealMatrix(oQMatrix);
            RealMatrix oRealRMatrix = new Array2DRowRealMatrix(oRMatrix);

            RealMatrix oIdentityMatrix = MatrixUtils.createRealIdentityMatrix(oRealQMatrix.getRowDimension());

            m_oBMatrix = ((new LUDecompositionImpl(oIdentityMatrix.subtract(oRealQMatrix)).getSolver().getInverse()).multiply(oRealRMatrix)).getData();

        }
    }

    /**
     * Get the probability of reaching a absorbing state from a transient state.
     *
     * @param transientState    the transient state.
     * @param absorbingTextPart the absorbing text part.
     * @param absorbingSetPart  the absorbing set part.
     * @return the probability of reaching reaching the absorbing state from the transient state.
     */
    public Double getProbability(RMPIMarkovState transientState, String absorbingTextPart, Set<String> absorbingSetPart) {
        generateAndSolveMarkovChain();

        return m_oBMatrix[transientState.getID()][getState(absorbingTextPart, absorbingSetPart, true).getID()];
    }

    /**
     * Get the probability of reaching a absorbing state from a transient state.
     *
     * @param transientTextPart the transient text part.
     * @param transientSetPart  the transient set part.
     * @param absorbingTextPart the absorbing text part.
     * @param absorbingSetPart  the absorbing set part.
     * @return the probability of reaching reaching the absorbing state from the transient state.
     */
    @Deprecated
    public Double getProbability(String transientTextPart, Set<String> transientSetPart, String absorbingTextPart, Set<String> absorbingSetPart) {

        return getProbability(getState(transientTextPart, transientSetPart, false), absorbingTextPart, absorbingSetPart);
    }

    /**
     * Add a transient state.
     *
     * @param textPart the text part.
     * @param setPart  the set part.
     */
    @Deprecated
    public void addTransientMarkovState(String textPart, Set<String> setPart) {
        addState(textPart, setPart, false);
    }

    /**
     * Add an absorbing state.
     *
     * @param textPart the text part.
     * @param setPart  the set part.
     */
    @Deprecated
    public void addAbsorbingMarkovState(String textPart, Set<String> setPart) {
        addState(textPart, setPart, true);
    }

    /**
     * Get a transient state.
     *
     * @param textPart the text part.
     * @param setPart  the set part.
     */
    @Deprecated
    public RMPIMarkovState getTransientStateByLabel(String textPart, Set<String> setPart) {
        return getState(textPart, setPart, false);
    }

    /**
     * Get an absorbing state.
     *
     * @param textPart the text part.
     * @param setPart  the set part.
     */
    @Deprecated
    public RMPIMarkovState getAbsorbingStateByLabel(String textPart, Set<String> setPart) {
        return getState(textPart, setPart, true);
    }


    /**
     * Add an absorbing transition.
     *
     * @param sourceTextPart the text part of the source state.
     * @param sourceSetPart  the set part of the source state.
     * @param desTextPart    the text part of the destination state.
     * @param desSetPart     the set part of the destination state.
     * @param probability    the transition probability.
     */
    @Deprecated
    public void addAbsorbingMarkovTransition(String sourceTextPart, HashSet<String> sourceSetPart, String desTextPart, HashSet<String> desSetPart, double probability) {
        RMPIMarkovState oSourceState = getState(sourceTextPart, sourceSetPart, false);
        RMPIMarkovState oDesState = getState(desTextPart, desSetPart, true);

        RMPIMarkovTransition oTransition = new RMPIMarkovTransition(oSourceState, oDesState, probability);
        oSourceState.addOutgoingTransition(oTransition);
    }

    /**
     * Add an transient transition.
     *
     * @param sourceTextPart the text part of the source state.
     * @param sourceSetPart  the set part of the source state.
     * @param desTextPart    the text part of the destination state.
     * @param desSetPart     the set part of the destination state.
     * @param probability    the transition probability.
     */
    @Deprecated
    public void addTransientMarkovTransition(String sourceTextPart, HashSet<String> sourceSetPart, String desTextPart, HashSet<String> desSetPart, double probability) {
        RMPIMarkovState oSourceState = getState(sourceTextPart, sourceSetPart, false);
        RMPIMarkovState oDesState = getState(desTextPart, desSetPart, false);

        RMPIMarkovTransition oTransition = new RMPIMarkovTransition(oSourceState, oDesState, probability);
        oSourceState.addOutgoingTransition(oTransition);
    }
}
